import torch
from torch import nn
from torch.nn import functional as F
from typing import Any, Dict, List, Tuple
from algorithms.SAM.image_encoder import ImageEncoderViT
from algorithms.SAM.mask_decoder import MaskDecoder
from algorithms.SAM.prompt_encoder import PromptEncoder


class Sam(nn.Module):
    mask_threshold: float = 0.0
    image_format: str = "RGB"

    def __init__(
            self,
            image_encoder: ImageEncoderViT,
            prompt_encoder: PromptEncoder,
            mask_decoder: MaskDecoder,
            pixel_mean: List[float] = [123.675, 116.28, 103.53],
            pixel_std: List[float] = [58.395, 57.12, 57.375],
    ) -> None:
        """
        SAM模型，根据image prompt和input prompt中进行预测。
        Args:
          image_encoder (ImageEncoderViT): The backbone used to encode the
            image into image embeddings that allow for efficient mask prediction.
          prompt_encoder (PromptEncoder): Encodes various types of input prompts.
          mask_decoder (MaskDecoder): Predicts masks from the image embeddings
            and encoded prompts.
          pixel_mean (list(float)): Mean values for normalizing pixels in the input image.
          pixel_std (list(float)): Std values for normalizing pixels in the input image.
        """
        super().__init__()
        self.image_encoder = image_encoder
        self.prompt_encoder = prompt_encoder
        self.mask_decoder = mask_decoder
        self.pixel_mean = torch.Tensor(pixel_mean).view(-1, 1, 1)
        self.pixel_std = torch.Tensor(pixel_std).view(-1, 1, 1)
        # self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False)
        # self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)
        # if "pixel_mean" in self._buffers:
        #     del self._buffers["pixel_mean"]
        #     self.register_buffer("pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False)
        # if "pixel_std" in self._buffers:
        #     del self._buffers["pixel_std"]
        #     self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)

    @property
    def device(self) -> Any:
        return self.pixel_mean.device

    @torch.no_grad()
    def forward(
            self,
            batched_input: List[Dict[str, Any]],
            multi_mask_output: bool,
    ) -> List[Dict[str, torch.Tensor]]:
        """
        根据提供的图像和prompts预测掩码，如果prompts为止，建议使用SamPredictor，而非直接调用该模型
        Args:
          batched_input (list(dict)): 每个dict包含以下内容。prompt可以不包含其中。
              'image': (torch.Tensor) 经过转换的图像数据，(3,H,W)
              'original_size': (tuple(int, int)) 转换前的原始图像大小，(H, W).
              'point_coords': (torch.Tensor) 针对当前图像的批量点Prompts，(B,N,2)，已经过转换，作为模型的输入。
              'point_labels': (torch.Tensor) 批量点Prompts对应的labels，(B,N)。
              'boxes': (torch.Tensor) 批量box输入，(Bx4)。已经过转换.
              'mask_inputs': (torch.Tensor) 批量的mask输入，(B,1,H,W)。
          multi_mask_output (bool): 是否预测多个无歧义的mask，还是仅返回一个mask。
        Returns:
          (list(dict)): 列表，其中每个元素都是一个包含以下键的字典。
              'masks': (torch.Tensor) 批量二进制mask预测，(B,C,H,W)，B是输入prompt的数量，C由mulimask_output决定，
                指的是一个mask中是否需要包含多个分类，(H,W)是原始图像的大小。
              'iou_predictions': (torch.Tensor) 模型针对每个mask质量的评估结果，(BxC)。
              'low_res_logits': (torch.Tensor) Low resolution logits with
                shape BxCxHxW, where H=W=256. Can be passed as mask input
                to subsequent iterations of prediction.
        """
        # 批量提取batched_input中的image，并堆叠成批量数据
        # preprocess对图像数据进行归一化，以及对于大小进行padding
        input_images = torch.stack([self.preprocess(x["image"]) for x in batched_input], dim=0)
        # 将批量图像数据传给image_encoder，进行embedding
        image_embeddings = self.image_encoder(input_images)

        outputs = []
        for image_record, curr_embedding in zip(batched_input, image_embeddings):
            if "point_coords" in image_record:
                points = (image_record["point_coords"], image_record["point_labels"])
            else:
                points = None
            sparse_embeddings, dense_embeddings = self.prompt_encoder(
                points=points,
                boxes=image_record.get("boxes", None),
                masks=image_record.get("mask_inputs", None),
            )
            low_res_masks, iou_predictions = self.mask_decoder(
                image_embeddings=curr_embedding.unsqueeze(0),
                image_pe=self.prompt_encoder.get_dense_pe(),
                sparse_prompt_embeddings=sparse_embeddings,
                dense_prompt_embeddings=dense_embeddings,
                multi_mask_output=multi_mask_output,
            )
            masks = self.postprocess_masks(
                low_res_masks,
                input_size=image_record["image"].shape[-2:],
                original_size=image_record["original_size"],
            )
            masks = masks > self.mask_threshold
            outputs.append(
                {
                    "masks": masks,
                    "iou_predictions": iou_predictions,
                    "low_res_logits": low_res_masks,
                }
            )
        return outputs

    def postprocess_masks(
            self,
            masks: torch.Tensor,
            input_size: Tuple[int, ...],
            original_size: Tuple[int, ...],
    ) -> torch.Tensor:
        """
        Remove padding and upscale masks to the original image size.

        Arguments:
          masks (torch.Tensor): Batched masks from the mask_decoder,
            in BxCxHxW format.
          input_size (tuple(int, int)): The size of the image input to the
            model, in (H, W) format. Used to remove padding.
          original_size (tuple(int, int)): The original size of the image
            before resizing for input to the model, in (H, W) format.

        Returns:
          (torch.Tensor): Batched masks in BxCxHxW format, where (H, W)
            is given by original_size.
        """
        masks = F.interpolate(
            masks,
            (self.image_encoder.img_size, self.image_encoder.img_size),
            mode="bilinear",
            align_corners=False,
        )
        masks = masks[..., : input_size[0], : input_size[1]]
        masks = F.interpolate(masks, original_size, mode="bilinear", align_corners=False)
        return masks

    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
        """Normalize pixel values and pad to a square input."""
        # Normalize colors
        self.pixel_mean = self.pixel_mean.to(x.device)
        self.pixel_std = self.pixel_std.to(x.device)
        x = (x - self.pixel_mean) / self.pixel_std

        # Pad
        h, w = x.shape[-2:]
        padh = self.image_encoder.img_size - h
        padw = self.image_encoder.img_size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x
